import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from _data_process import cumulative_to_monthly
from _data_process import monthly_to_quarterly
from _data_process import calculate_yoy
# from _data_plot import plot_line_chart

import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.ticker import MaxNLocator


def plot_line_chart(fig_width, fig_height, x_data, y_data, title, x_label, y_label):
    """
    创建并返回一个折线图
    
    参数:
        fig_width: 图形宽度
        fig_height: 图形高度
        x_data: x轴数据（通常是日期）
        y_data: y轴数据
        title: 图表标题
        x_label: x轴标签
        y_label: y轴标签
    
    返回:
        fig: 生成的matplotlib图形对象
    """
    # 创建图形和坐标轴
    fig, ax = plt.subplots(figsize=(fig_width, fig_height))
    plt.rcParams["font.family"] = ["SimHei"]
    plt.rcParams["axes.unicode_minus"] = False
    # 绘制折线图
    ax.plot(x_data, y_data, marker='o', linestyle='-', color='b', markersize=4, linewidth=1.5)
    
    # 设置标题和标签
    ax.set_title(title, fontsize=14, pad=15)
    ax.set_xlabel(x_label, fontsize=12, labelpad=10)
    ax.set_ylabel(y_label, fontsize=12, labelpad=10)
    
    # 美化x轴日期显示（如果x是日期类型）
    if hasattr(x_data, 'dtype') and 'datetime64' in str(x_data.dtype):
        ax.xaxis.set_major_locator(mdates.AutoDateLocator())
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
        plt.xticks(rotation=45, ha='right')  # 旋转日期标签避免重叠
    
    # 设置y轴刻度为整数（如果数据是整数类型）
    if all(isinstance(y, (int, np.integer)) for y in y_data):
        ax.yaxis.set_major_locator(MaxNLocator(integer=True))
    
    # 添加网格线
    ax.grid(True, linestyle='--', alpha=0.7)
    
    # 调整布局
    plt.tight_layout()
    
    return fig
series_data = pd.read_excel(f"data/全社会用电量当月累计值.xlsx", index_col=0)
x_cum = series_data.index
y_cum = series_data.iloc[:,0].values
# 绘制折线图（宽度10，高度6）
fig_cum = plot_line_chart(10, 6, 
                        x_cum, y_cum, 
                        "全社会用电量", "日期", "累计值")

series_data_monthly = cumulative_to_monthly(series_data)
data = series_data_monthly.iloc[:,0]
x_monthly = data.index
y_monthly = data.values

# 绘制折线图（宽度10，高度6）
fig_monthly = plot_line_chart(10, 6, 
                        x_monthly, y_monthly, "全社会用电量", "日期", "当月值")

series_data_quarterly = monthly_to_quarterly(series_data_monthly)
data = series_data_quarterly.iloc[:,0]
x = data.index
y = data.values
# 绘制折线图（宽度10，高度6）
fig_quarterly = plot_line_chart(10, 6, 
                        x, y, "全社会用电量", "日期", "季度值")


series_data_quarterly_yoy = calculate_yoy(series_data_quarterly)

data = series_data_quarterly_yoy.iloc[:,0]
x_yoy = data.index
y_yoy = data.values
# 绘制折线图（宽度10，高度6）
fig_quarterly_yoy = plot_line_chart(10, 6, 
                        x_yoy, y_yoy, "全社会用电量", "日期", "季度同比值")

# 显示图像
plt.show()
# 如果需要保存图像，可以使用
# fig.savefig("line_chart.png", bbox_inches='tight')